{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define a model that serves as an example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Export a PyTorch model to an ONNX model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.5.0\n"
     ]
    }
   ],
   "source": [
    "import torchvision\n",
    "import torch\n",
    "import os\n",
    "import time\n",
    "from PIL import Image\n",
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "print(torch.__version__)\n",
    "\n",
    "onnx_model_path = 'resnet50.onnx'\n",
    "\n",
    "input_name = ['input']\n",
    "output_name = ['output']\n",
    "input = torch.randn(1, 3, 224, 224).cuda()\n",
    "model = torchvision.models.resnet50(pretrained=True).cuda()\n",
    "torch.onnx.export(model, input, onnx_model_path, input_names=input_name, output_names=output_name, verbose=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### check onnx model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import onnx\n",
    "# onnx_model_path = 'resnet50.onnx'\n",
    "\n",
    "# test = onnx.load(onnx_model_path)\n",
    "# onnx.checker.check_model(test)\n",
    "# print(\"==> Passed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### convert an onnx model to a TensorRT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorrt as trt\n",
    "\n",
    "TRT_LOGGER = trt.Logger(trt.Logger.WARNING)\n",
    "EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)\n",
    "\n",
    "def get_engine(max_batch_size=1, onnx_file_path=\"\", engine_file_path=\"\", \\\n",
    "               fp16_mode=False, int8_mode=False, save_engine=True\n",
    "               ):\n",
    "    \"\"\"Attempts to load a serialized engine if available, otherwise builds a new TensorRT engine and saves it.\"\"\"\n",
    "\n",
    "    def build_engine(max_batch_size, save_engine):\n",
    "        \"\"\"Takes an ONNX file and creates a TensorRT engine to run inference with\"\"\"\n",
    "        with trt.Builder(TRT_LOGGER) as builder, \\\n",
    "            builder.create_network(EXPLICIT_BATCH) as network, \\\n",
    "            trt.OnnxParser(network, TRT_LOGGER) as parser:\n",
    "            \n",
    "            # 设置 builder 参数\n",
    "            builder.max_workspace_size = 1 << 30  # Your workspace size\n",
    "            builder.max_batch_size = max_batch_size\n",
    "            builder.fp16_mode = fp16_mode  # Default: False\n",
    "            builder.int8_mode = int8_mode  # Default: False\n",
    "            if int8_mode:\n",
    "                raise NotImplementedError\n",
    "\n",
    "            # 解析模型\n",
    "            if not os.path.exists(onnx_file_path):\n",
    "                quit('ONNX file {} not found'.format(onnx_file_path))\n",
    "            print('Loading ONNX file from path {}...'.format(onnx_file_path))\n",
    "            with open(onnx_file_path, 'rb') as model:\n",
    "                print('Beginning ONNX file parsing')\n",
    "                parser.parse(model.read())\n",
    "            print('Completed parsing of ONNX file')\n",
    "            print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))\n",
    "\n",
    "            # 构建推理engine\n",
    "            engine = builder.build_cuda_engine(network)\n",
    "            print(\"Completed creating Engine\")\n",
    "\n",
    "            if save_engine:\n",
    "                with open(engine_file_path, \"wb\") as f:\n",
    "                    f.write(engine.serialize())\n",
    "            return engine\n",
    "\n",
    "    if os.path.exists(engine_file_path):\n",
    "        # If a serialized engine exists, load it instead of building a new one.\n",
    "        print(\"Reading engine from file {}\".format(engine_file_path))\n",
    "        with open(engine_file_path, \"rb\") as f, trt.Runtime(TRT_LOGGER) as runtime:\n",
    "            return runtime.deserialize_cuda_engine(f.read())\n",
    "    else:\n",
    "        return build_engine(max_batch_size, save_engine)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading ONNX file from path resnet50.onnx...\n",
      "Beginning ONNX file parsing\n",
      "Completed parsing of ONNX file\n",
      "Building an engine from file resnet50.onnx; this may take a while...\n",
      "Completed creating Engine\n"
     ]
    }
   ],
   "source": [
    "max_batch_size = 1\n",
    "# These two modes are dependent on hardwares\n",
    "fp16_mode = True\n",
    "int8_mode = False\n",
    "trt_engine_path = './model_fp16_{}_int8_{}.trt'.format(fp16_mode, int8_mode)\n",
    "# Build an engine\n",
    "engine = get_engine(max_batch_size, onnx_model_path, trt_engine_path, fp16_mode, int8_mode)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Do inference with Pytorch & onnx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_img_np_nchw(filename):\n",
    "    image = cv2.imread(filename)\n",
    "    image_cv = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    image_cv = cv2.resize(image_cv, (224, 224))\n",
    "    miu = np.array([0.485, 0.456, 0.406])\n",
    "    std = np.array([0.229, 0.224, 0.225])\n",
    "    img_np = np.array(image_cv, dtype=float) / 255.\n",
    "    r = (img_np[:, :, 0] - miu[0]) / std[0]\n",
    "    g = (img_np[:, :, 1] - miu[1]) / std[1]\n",
    "    b = (img_np[:, :, 2] - miu[2]) / std[2]\n",
    "    img_np_t = np.array([r, g, b])\n",
    "    img_np_nchw = np.expand_dims(img_np_t, axis=0)\n",
    "    return img_np_nchw\n",
    "\n",
    "filename = 'cat.jpeg'\n",
    "img_np_nchw = get_img_np_nchw(filename)\n",
    "img_np_nchw = img_np_nchw.astype(dtype=np.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### (1) Pytorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005634307861328125\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005437135696411133\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005404472351074219\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.00537562370300293\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.00644683837890625\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005829811096191406\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005459308624267578\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005388736724853516\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.0053789615631103516\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005362272262573242\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005383014678955078\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005387067794799805\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005362033843994141\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.0053806304931640625\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.0053920745849609375\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.00594639778137207\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.006297111511230469\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005574226379394531\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.0054318904876708984\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005314826965332031\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.0052835941314697266\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005488157272338867\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005422115325927734\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005710124969482422\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005519866943359375\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005429267883300781\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005713224411010742\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.006940364837646484\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005455493927001953\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005464792251586914\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.0052814483642578125\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005317211151123047\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005405426025390625\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.0053501129150390625\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005712032318115234\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.0055103302001953125\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005685329437255859\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.0055828094482421875\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.006972789764404297\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005492687225341797\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005454301834106445\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.00539398193359375\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005368947982788086\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005376100540161133\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005383014678955078\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005357265472412109\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.005365610122680664\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.00535273551940918\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.00537419319152832\n",
      "Pytorch ok!\n",
      "Inference time with the PyTorch model: 0.006824970245361328\n"
     ]
    }
   ],
   "source": [
    "model = torchvision.models.resnet50(pretrained=True).cuda()\n",
    "resnet_model = model.eval()\n",
    "\n",
    "for i in range(50):\n",
    "    input_for_torch = torch.from_numpy(img_np_nchw).cuda()\n",
    "    t3 = time.time()\n",
    "    feat_2= resnet_model(input_for_torch)\n",
    "    t4 = time.time()\n",
    "    feat_2 = feat_2.cpu().data.numpy()\n",
    "    print('Pytorch ok!')\n",
    "\n",
    "    print(\"Inference time with the PyTorch model: {}\".format(t4-t3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### (2) TensorRT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pycuda.driver as cuda\n",
    "import pycuda.autoinit\n",
    "\n",
    "\n",
    "class HostDeviceMem(object):\n",
    "    def __init__(self, host_mem, device_mem):\n",
    "        \"\"\"Within this context, host_mom means the cpu memory and device means the GPU memory\n",
    "        \"\"\"\n",
    "        self.host = host_mem\n",
    "        self.device = device_mem\n",
    "\n",
    "    def __str__(self):\n",
    "        return \"Host:\\n\" + str(self.host) + \"\\nDevice:\\n\" + str(self.device)\n",
    "\n",
    "    def __repr__(self):\n",
    "        return self.__str__()\n",
    "    \n",
    "def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):\n",
    "    # Transfer data from CPU to the GPU.\n",
    "    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]\n",
    "    # Run inference.\n",
    "    context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)\n",
    "    # Transfer predictions back from the GPU.\n",
    "    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]\n",
    "    # Synchronize the stream\n",
    "    stream.synchronize()\n",
    "    # Return only the host outputs.\n",
    "    return [out.host for out in outputs]\n",
    "\n",
    "\n",
    "def postprocess_the_outputs(h_outputs, shape_of_output):\n",
    "    h_outputs = h_outputs.reshape(*shape_of_output)\n",
    "    return h_outputs\n",
    "\n",
    "def allocate_buffers(engine):\n",
    "    inputs = []\n",
    "    outputs = []\n",
    "    bindings = []\n",
    "    stream = cuda.Stream()\n",
    "    for binding in engine:\n",
    "        size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size\n",
    "        dtype = trt.nptype(engine.get_binding_dtype(binding))\n",
    "        # Allocate host and device buffers\n",
    "        host_mem = cuda.pagelocked_empty(size, dtype)\n",
    "        device_mem = cuda.mem_alloc(host_mem.nbytes)\n",
    "        # Append the device buffer to device bindings.\n",
    "        bindings.append(int(device_mem))\n",
    "        # Append to the appropriate list.\n",
    "        if engine.binding_is_input(binding):\n",
    "            inputs.append(HostDeviceMem(host_mem, device_mem))\n",
    "        else:\n",
    "            outputs.append(HostDeviceMem(host_mem, device_mem))\n",
    "    return inputs, outputs, bindings, stream"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the context for this engine\n",
    "context = engine.create_execution_context()\n",
    "# # Allocate buffers for input and output\n",
    "inputs, outputs, bindings, stream = allocate_buffers(engine) # input, output: host # bindings\n",
    "\n",
    "# Do inference\n",
    "shape_of_output = (max_batch_size, 1000)\n",
    "# Load data to the buffer\n",
    "inputs[0].host = img_np_nchw.reshape(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0014615058898925781\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0012578964233398438\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011913776397705078\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011141300201416016\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.001186370849609375\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0012629032135009766\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011968612670898438\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0012068748474121094\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0021479129791259766\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0015397071838378906\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0012843608856201172\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.00125885009765625\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.001178741455078125\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.001127481460571289\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011947154998779297\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011172294616699219\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0017178058624267578\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011878013610839844\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011892318725585938\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0012183189392089844\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0018963813781738281\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0014078617095947266\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0015726089477539062\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0012638568878173828\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011909008026123047\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011920928955078125\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011878013610839844\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011947154998779297\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0012028217315673828\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.001195669174194336\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011982917785644531\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011935234069824219\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0018167495727539062\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0013184547424316406\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011913776397705078\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0012810230255126953\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011937618255615234\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011761188507080078\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011899471282958984\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011954307556152344\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0017168521881103516\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0012369155883789062\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0012087821960449219\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0012195110321044922\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0012402534484863281\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0012874603271484375\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0011205673217773438\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0014722347259521484\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0012621879577636719\n",
      "TensorRT ok\n",
      "Inference time with the TensorRT engine: 0.0012488365173339844\n",
      "MSE Error = 5.8050314692081884e-05\n",
      "All completed!\n"
     ]
    }
   ],
   "source": [
    "for i in range(50):\n",
    "    t1 = time.time()\n",
    "    trt_outputs = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream) # numpy data\n",
    "    t2 = time.time()\n",
    "    feat = postprocess_the_outputs(trt_outputs[0], shape_of_output)\n",
    "\n",
    "    print('TensorRT ok')\n",
    "    print(\"Inference time with the TensorRT engine: {}\".format(t2-t1))\n",
    "    \n",
    "print('MSE Error = {}'.format(np.mean((feat - feat_2)**2)))\n",
    "\n",
    "print('All completed!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Experiments\n",
    "\n",
    "##### Test Device: GTX 2080Ti        \n",
    "##### Network:Resnet50                 \n",
    "##### Inputsize:224x224\n",
    "\n",
    "##### Pytorch Inference: 0.0053s\n",
    "\n",
    "##### TensorRT with FP32: 0.0027s  \n",
    "\n",
    "##### TensorRT with FP16: 0.0017s\n",
    "\n",
    "##### TensorRT with FP16(change workspace from 1 >> 20 to 1>> 32): 0.0012s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
